package cn.xiaopengstack.config;

import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * @author jiangyangang
 */
@Configuration
@EnableConfigurationProperties(SpringAiProperties.class)
public class RAGEmbeddingConfig {

    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }


    @Bean("ollamaSimpleVectorStore")
    public SimpleVectorStore vectorStore(SpringAiProperties properties, OllamaApi ollamaApi) {
        OllamaEmbeddingModel embeddingModel = OllamaEmbeddingModel
                .builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder().model(properties.getEmbeddingOptionsModel()).build())
                .build();
        return SimpleVectorStore.builder(embeddingModel).build();
    }

    @Bean("ollamaPgVectorStore")
    public PgVectorStore pgVectorStore(SpringAiProperties properties, OllamaApi ollamaApi, JdbcTemplate jdbcTemplate) {
        OllamaEmbeddingModel embeddingModel = OllamaEmbeddingModel
                .builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder().model(properties.getEmbeddingOptionsModel()).build())
                .build();
        return PgVectorStore.builder(jdbcTemplate, embeddingModel)
                .vectorTableName(properties.getPostgresTableName())
                .dimensions(properties.getEmbeddingDimensions())
                .build();
    }
}
